﻿#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
图像处理器模块

提供图像加载、保存和处理功能
"""

import cv2
import numpy as np
from PIL import Image
import os
import tempfile
import shutil
import traceback
import uuid

class ImageProcessor:
    """图像处理类，提供图像处理核心功能"""
    
    @staticmethod
    def load_image(file_path):
        """加载图像文件
        
        参数:
            file_path (str): 图像文件路径
            
        返回:
            numpy.ndarray: 加载的图像数据(RGB格式)，加载失败返回None
        """
        try:
            # 使用PIL加载图像（支持中文路径）
            pil_image = Image.open(file_path)
            # 转换为RGB格式
            image = np.array(pil_image.convert('RGB'))
            return image
        except Exception as e:
            print(f"加载图像时出错: {str(e)}")
            traceback.print_exc()
            return None
            
    @staticmethod
    def save_image(image, file_path, format="PNG"):
        """保存图像到文件
        
        参数:
            image (numpy.ndarray): 图像数据
            file_path (str): 保存路径
            format (str): 图像格式，可选"PNG", "JPEG", "BMP"等
            
        返回:
            bool: 保存成功返回True，否则False
        """
        try:
            print(f"开始保存图像: 路径={file_path}, 格式={format}")
            print(f"图像信息: 形状={image.shape}, 类型={image.dtype}")
            print(f"图像左上角像素值: {image[0, 0] if image.size > 0 else 'N/A'}")
            
            # 检查并规范化文件路径
            file_path = os.path.abspath(file_path)
            print(f"规范化后的文件路径: {file_path}")
            
            # 检查文件名是否有效
            file_dir = os.path.dirname(file_path)
            file_name = os.path.basename(file_path)
            
            print(f"文件目录: {file_dir}")
            print(f"文件名: {file_name}")
            
            # 检查文件名是否包含非法字符
            invalid_chars = '<>:"/\\|?*'
            has_invalid_chars = False
            for c in file_name:
                if c in invalid_chars:
                    has_invalid_chars = True
                    print(f"文件名包含非法字符: '{c}'")
            
            if has_invalid_chars:
                # 尝试修复文件名
                safe_name = ""
                for c in file_name:
                    if c in invalid_chars:
                        safe_name += "_"  # 替换非法字符为下划线
                    else:
                        safe_name += c
                print(f"修复后的文件名: {safe_name}")
                file_path = os.path.join(file_dir, safe_name)
                print(f"修复后的文件路径: {file_path}")
            
            # 确保目录存在
            if file_dir and not os.path.exists(file_dir):
                print(f"目录不存在，尝试创建: {file_dir}")
                os.makedirs(file_dir, exist_ok=True)
                print(f"目录创建成功: {file_dir}")
            
            # 检查图像是否有效
            if image is None or image.size == 0:
                print("错误: 图像数据无效")
                return False
            
            # 处理格式兼容性 - JPG应该使用JPEG标识符
            save_format = format.upper()
            if save_format == 'JPG':
                save_format = 'JPEG'
                print("格式'JPG'已转换为PIL支持的'JPEG'格式")
                
            # 直接使用PIL保存图像，避免颜色转换
            print("直接使用PIL保存图像，避免颜色空间转换")
            try:
                # 创建PIL图像对象
                pil_image = Image.fromarray(image)
                print(f"PIL图像信息: 尺寸={pil_image.size}, 模式={pil_image.mode}")
                
                # 直接保存
                pil_image.save(file_path, save_format)
                print(f"已保存图像: {file_path}")
                
                # 验证保存的图像
                verify_img = np.array(Image.open(file_path))
                print(f"验证保存的图像: 形状={verify_img.shape}, 类型={verify_img.dtype}")
                print(f"保存图像左上角像素值: {verify_img[0, 0]}")
                print(f"原图与保存图像像素值是否一致: {np.array_equal(image[0, 0], verify_img[0, 0])}")
                
                return True
            except Exception as e:
                print(f"直接保存失败: {e}，尝试使用临时文件中转")
                
                # 使用临时文件中转的方式来修正颜色问题
                temp_dir = tempfile.gettempdir()
                temp_filename = f"temp_img_{uuid.uuid4().hex}.png"
                temp_path = os.path.join(temp_dir, temp_filename)
                
                print(f"使用临时文件中转方式保存: {temp_path}")
                
                # 假设输入的图像是RGB格式，保存为临时文件
                pil_image = Image.fromarray(image)
                pil_image.save(temp_path, "PNG")
                print(f"已保存临时文件: {temp_path}")
                
                # 再次打开这个文件并保存为最终格式
                final_img = np.array(Image.open(temp_path))
                final_pil = Image.fromarray(final_img)
                
                try:
                    final_pil.save(file_path, save_format)
                    print(f"已保存最终图像: {file_path}")
                except PermissionError as pe:
                    print(f"权限错误: {str(pe)}")
                    # 尝试使用替代路径
                    alt_path = os.path.join(os.path.expanduser("~"), os.path.basename(file_path))
                    final_pil.save(alt_path, save_format)
                    print(f"已保存到替代路径: {alt_path}")
                    file_path = alt_path
                
                # 删除临时文件
                try:
                    os.remove(temp_path)
                    print(f"已删除临时文件: {temp_path}")
                except Exception as e:
                    print(f"无法删除临时文件: {e}")
                
                # 验证最终保存的图像
                try:
                    final_img_check = np.array(Image.open(file_path))
                    print(f"验证保存的图像: 形状={final_img_check.shape}, 类型={final_img_check.dtype}")
                    print(f"保存图像左上角像素值: {final_img_check[0, 0]}")
                    print(f"原图与保存图像像素值是否一致: {np.array_equal(image[0, 0], final_img_check[0, 0])}")
                except Exception as e:
                    print(f"验证保存的图像时出错: {e}")
                
                return True
                
        except Exception as e:
            print(f"保存图像时出错: {str(e)}")
            print(f"错误详情: 文件路径={file_path}, 图像尺寸={image.shape if hasattr(image, 'shape') else '未知'}")
            traceback.print_exc()
            return False
    
    @staticmethod
    def remove_background(image, target_color, threshold):
        """移除图像背景
        
        参数:
            image: 输入图像
            target_color: 待滤除的颜色列表 [R, G, B]
            threshold: 颜色差异阈值
            
        返回:
            处理后的图像（带透明通道）
        """
        if image is None:
            return None
        
        # 创建图像副本
        result = image.copy()

        # 将目标颜色从RGB转换为BGR格式
        bgr_target_color = [target_color[2], target_color[1], target_color[0]]
        print(f"目标颜色RGB: {target_color} -> BGR: {bgr_target_color}")
        
        # 计算颜色差异
        diff = np.abs(result - np.array(target_color))
        mask = np.sum(diff, axis=2) <= threshold
        
        # 创建RGBA图像
        rgba = cv2.cvtColor(result, cv2.COLOR_BGR2BGRA)
        
        # 只修改匹配像素的alpha通道为0（完全透明）
        rgba[mask, 3] = 0
        
        return rgba
    
    @staticmethod
    def detect_area(image, x, y, sensitivity):
        """根据点击位置智能检测区域
        
        参数:
            image (numpy.ndarray): 原始图像
            x (int): 点击位置X坐标
            y (int): 点击位置Y坐标
            sensitivity (int): 像素差敏感度
            
        返回:
            tuple: 检测到的区域 (x, y, width, height) 或None
        """
        if image is None or not (0 <= x < image.shape[1] and 0 <= y < image.shape[0]):
            return None
        
        # 转换为灰度图
        gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
        
        # 应用高斯模糊减少噪声
        blurred = cv2.GaussianBlur(gray, (5, 5), 0)
        
        # 计算边缘
        edges = cv2.Canny(blurred, 50, 150)
        
        # 使用漫水填充算法寻找区域
        h, w = gray.shape
        mask = np.zeros((h+2, w+2), np.uint8)
        
        # 设定漫水填充参数
        flood_fill_flags = 4  # 4邻域连通
        flood_fill_flags |= cv2.FLOODFILL_FIXED_RANGE
        flood_fill_flags |= (255 << 8)  # 填充值
        
        # 执行漫水填充
        _, _, _, rect = cv2.floodFill(
            edges.copy(), mask, (x, y), 255,
            (sensitivity,)*3, (sensitivity,)*3,
            flood_fill_flags
        )
        
        # 如果找到的区域太小，返回None
        if rect[2] < 10 or rect[3] < 10:
            return None
        
        return rect

    @staticmethod
    def detect_multiple_areas(image, mouse_x, mouse_y, sensitivity=30):
        """检测图像中的多个可能区域，确保包含鼠标位置
        
        参数:
            image: 输入图像
            mouse_x: 鼠标X坐标
            mouse_y: 鼠标Y坐标
            sensitivity: 检测敏感度
            
        返回:
            list: 建议区域列表，每个区域为 (x, y, w, h)
        """
        # 转换为灰度图
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        
        # 高斯模糊
        blurred = cv2.GaussianBlur(gray, (5, 5), 0)
        
        # Canny边缘检测
        edges = cv2.Canny(blurred, 40, 120)
        
        # 查找轮廓
        contours, _ = cv2.findContours(edges, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
        
        # 图像尺寸
        height, width = image.shape[:2]
        max_area = width * height * 0.95  # 最大面积为图像的95%
        min_area = 100  # 最小面积
        
        # 存储有效的矩形
        valid_rects = []
        
        # 遍历轮廓
        for contour in contours:
            # 计算轮廓面积
            area = cv2.contourArea(contour)
            
            # 面积过滤
            if area < min_area or area > max_area:
                continue
            
            # 获取外接矩形
            x, y, w, h = cv2.boundingRect(contour)
            
            # 检查矩形是否包含鼠标位置
            if x <= mouse_x <= x + w and y <= mouse_y <= y + h:
                # 确保矩形不会太小
                if w >= 5 and h >= 5:
                    valid_rects.append((x, y, w, h, area))  # 添加面积信息用于排序
        
        if not valid_rects:
            return []
        
        # 按面积排序，选择最小的区域
        valid_rects.sort(key=lambda x: x[4])  # 按面积排序
        smallest_rect = valid_rects[0][:4]  # 只返回坐标信息
        
        return [smallest_rect]  # 返回单个区域